import torch
import torch.nn as nn


class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.deconv1 = nn.ConvTranspose2d(in_channels=512,out_channels=256,kernel_size=2,stride=2)
        self.deconv2 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=2, stride=2)
        self.deconv3 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=2, stride=2)
        self.deconv4 = nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=2, stride=2)
        self.deconv5 = nn.ConvTranspose2d(in_channels=64, out_channels=1, kernel_size=2, stride=2)

        self.conv1=nn.Conv2d(in_channels=512,out_channels=256,kernel_size=3,padding=1)
        self.conv2 = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=1)

    def forward(self, feats):  ## x is a list

        # for idx ,feats in enumerate(feats[:-1]):##如何逆序遍历？

        up1 = self.deconv1(feats[4])    ##  12->24
        out1 = torch.cat((up1,feats[3]),dim=1)
        out1=self.conv1(out1)            # out1 [256,24,24]

        up2=self.deconv2(out1)          ##  24->48
        out2=torch.cat((up2,feats[2]),dim=1)
        out2=self.conv2(out2)           ##out2 [128,48,48]

        up3=self.deconv3(out2)          ## 48->96
        out3=torch.cat((up3,feats[1]),dim=1)
        out3=self.conv3(out3)                   ## out3 [64,96,96]

        up4=self.deconv4(out3)           ##96->192
        out4=torch.cat((up4,feats[0]),dim=1)
        out4=self.conv4(out4)                  ## out4 [64, 192,192]

        out = self.deconv5(out4)                  ##   [1,1,384,384]
        out = nn.Sigmoid()(out)

        return out


